import functools
import heapq

from google.appengine.ext import db

from lib import cache
from lib import models


class Error(Exception): pass
class DeadlockError(Error): pass


class State(object):

  def __init__(self, version_str, race_str, map_str):
    self.event_history = []
    self.time = 0.0
    self.current_supply = 0
    self.max_supply = 0
    self.queue = []
    self.future = []
    self.workers_by_group = {}
    self.building = set()
    self.building_by_source = {}
    self.objs_by_buildable = {}
    self.worker_groups = []

    (version_query, race_query, map_query) = cache.QueryMulti([
        db.Query(models.Version)
          .filter('version =', version_str),
        db.Query(models.Race)
          .filter('name =', race_str),
        db.Query(models.Map)
          .filter('name =', map_str),
    ])

    version = list(version_query)[0]
    race = list(race_query)[0]
    map = list(map_query)[0]
    self.minerals = map.starting_minerals
    self.gas = map.starting_gas
    (buildables_query, starting_state, base_query) = cache.QueryMulti([
        db.Query(models.Buildable)
          .filter('version =', version)
          .filter('race =', race),
        db.Query(models.StartingObject)
          .filter('map =', map)
          .filter('race =', race),
        db.Query(models.Base)
          .filter('map =', map),
    ])

    self.buildable_by_name = dict((x.name, x) for x in buildables_query)

    worker_group_queries = []
    for base in base_query:
      worker_group_queries.append(db.Query(models.WorkerGroup)
                                    .filter('base =', base))
    for worker_groups in cache.QueryMulti(worker_group_queries):
      self.worker_groups.extend(worker_groups)

    for starting_obj in starting_state:
      obj = self.CreateObject(starting_obj.buildable)
      self.current_supply += obj.buildable.cost_supply
      self.CompleteBuild(obj)

    self.AddEvent(models.EventType.GAME_START, None)

  def CreateObject(self, buildable):
    return models.Object(buildable=buildable)

  def AddEvent(self, type, obj, value=None):
    event = models.Event(type=type,
                         time=self.time,
                         current_supply=self.current_supply,
                         max_supply=self.max_supply,
                         minerals=self.minerals,
                         gas=self.gas)
    event.obj = obj
    if value:
      event.value = value
    self.event_history.append(event)

  def AddError(self, buildable, message):
    obj = self.CreateObject(buildable)
    self.AddEvent(models.EventType.ERROR, obj, message)
    raise DeadlockError(message)

  def AddFuture(self, offset, callback):
    heapq.heappush(self.future, (self.time + offset, callback))

  def StartBuild(self, buildable):
    obj = self.CreateObject(buildable)
    self.current_supply += buildable.cost_supply
    self.minerals -= buildable.cost_minerals
    self.gas -= buildable.cost_gas
    self.AddEvent(models.EventType.BUILD_START, obj)
    self.building.add(id(obj))
    if buildable.built_from:
      self.building_by_source.setdefault(buildable.built_from.key(), set()).add(obj)
    self.AddFuture(buildable.build_seconds,
                   functools.partial(self.CompleteBuild, obj))

  def CompleteBuild(self, obj):
    self.max_supply += obj.buildable.provides_supply
    self.AddEvent(models.EventType.BUILD_COMPLETE, obj)
    self.building.discard(id(obj))
    if obj.buildable.built_from:
      self.building_by_source.get(obj.buildable.built_from.key(), set()).discard(obj)
    self.objs_by_buildable.setdefault(obj.buildable.key(), set()).add(obj)
    if obj.buildable.worker:
      self.AssignWorker(obj, 'minerals')
    self.Dequeue()

  def AssignWorker(self, worker, resource_type):
    resource_type = resource_type.lower()
    for group in self.worker_groups:
      if (getattr(group, '%s_returned' % resource_type) and
          len(self.workers_by_group.get(group, [])) < group.max_workers):
        self.AddEvent(models.EventType.ASSIGN_WORKER, worker, resource_type)
        self.workers_by_group.setdefault(group, []).append(worker)
        self.AddFuture(group.resource_return_time,
                       functools.partial(self.ResourceReturn, worker, group))
        return
    assert False, 'Too many workers'

  def ResourceReturn(self, worker, worker_group):
    self.minerals += worker_group.minerals_returned
    self.gas += worker_group.gas_returned
    self.AddEvent(models.EventType.RESOURCE_RETURN, worker)
    self.AddFuture(worker_group.resource_return_time,
                   functools.partial(self.ResourceReturn, worker, worker_group))
    self.Dequeue()

  def Enqueue(self, name):
    self.queue.append(self.buildable_by_name[name])
    self.Dequeue()

  def Dequeue(self):
    """Dequeue and start build(s) if possible."""
    for buildable in list(self.queue):
      if (self.max_supply - self.current_supply) < buildable.cost_supply:
        if not self.building:
          self.AddError(buildable, 'Supply blocked')
        return
      if (self.minerals < buildable.cost_minerals or
          self.gas < buildable.cost_gas):
        return
      if buildable.built_from:
        if (len(self.objs_by_buildable.get(buildable.built_from.key(), [])) -
            len(self.building_by_source.get(buildable.built_from.key(), []))) < 1:
          if not self.building:
            self.AddError(buildable, 'Requires ' + buildable.built_from.name)
          return
      if buildable.requires:
        for requirement in buildable.requires:
          if requirement not in self.objs_by_buildable:
            if not self.building:
              self.AddError(buildable, 'Requires ' + models.Buildable.get(requirement).name)
            return
      self.queue.pop(0)
      self.StartBuild(buildable)

  def Run(self):
    while self.queue or self.building:
      (new_time, callback) = heapq.heappop(self.future)
      self.time = new_time
      callback()
